from flask import Flask, request, jsonify, Response, send_file, render_template
from flask_cors import CORS
import json
from typing import Generator
from config import config
from retriever import retriever
from backend import chat_manager
from database import db_manager

app = Flask(__name__)
CORS(app, resources={
    r"/api/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization"]
    }
})

@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/test_frontend.html')
def test_frontend():
    """测试页面"""
    return send_file('test_frontend.html')

@app.route('/api/health')
def health_check():
    """健康检查"""
    try:
        # 检查数据库连接
        doc_count = db_manager.get_document_count()
        
        # 检查检索器状态
        retriever_ready = retriever.initialized
        
        return jsonify({
            "status": "healthy",
            "database": {
                "connected": True,
                "document_count": doc_count
            },
            "retriever": {
                "initialized": retriever_ready
            },
            "config": {
                "max_history": config.MAX_HISTORY,
                "stream_enabled": config.STREAM_ENABLED
            }
        })
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e)
        }), 500

@app.route('/api/chat', methods=['POST'])
def chat():
    """聊天接口"""
    try:
        data = request.get_json()
        query = data.get('query', '').strip()
        search_type = data.get('search_type', 'hybrid')
        article_title = data.get('article_title', '').strip()
        stream = data.get('stream', config.STREAM_ENABLED)
        
        if not query:
            return jsonify({"error": "查询不能为空"}), 400
        
        if not article_title:
            return jsonify({"error": "文章标题不能为空"}), 400
        
        if stream:
            return Response(
                chat_stream_generator(query, search_type, article_title),
                mimetype='text/event-stream',
                headers={
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Access-Control-Allow-Origin': '*',
                    'Access-Control-Allow-Headers': 'Content-Type'
                }
            )
        else:
            # 非流式响应
            print(f"🔍 开始非流式聊天处理...")
            response_generator = chat_manager.generate_response(query, search_type, article_title, stream=False)
            response = next(response_generator)
            
            # 获取检索结果
            search_results = retriever.search(query, search_type, article_title)
            formatted_results = retriever.format_search_results(search_results)
            
            return jsonify({
                "response": response,
                "search_results": formatted_results,
                "conversation_summary": chat_manager.get_conversation_summary()
            })
            
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def chat_stream_generator(query: str, search_type: str, article_title: str) -> Generator[str, None, None]:
    """流式聊天生成器"""
    try:
        print(f"🔍 开始流式聊天处理...")
        
        # 获取检索结果
        search_results = retriever.search(query, search_type, article_title)
        formatted_results = retriever.format_search_results(search_results)
        
        # 发送检索结果
        yield f"data: {json.dumps({'search_results': formatted_results})}\n\n"
        
        # 生成流式回复
        print(f"🚀 调用chat_manager.generate_response...")
        for chunk in chat_manager.generate_response(query, search_type, article_title, stream=True):
            yield chunk
        
        # 发送对话摘要
        summary = chat_manager.get_conversation_summary()
        yield f"data: {json.dumps({'summary': summary})}\n\n"
        
    except Exception as e:
        print(f"❌ 流式生成异常: {e}")
        error_msg = f"流式生成异常: {str(e)}"
        yield f"data: {json.dumps({'error': error_msg})}\n\n"

@app.route('/api/search', methods=['POST'])
def search():
    """搜索接口"""
    try:
        data = request.get_json()
        query = data.get('query', '').strip()
        search_type = data.get('search_type', 'hybrid')
        article_title = data.get('article_title', '').strip()
        k = data.get('k', config.TOP_K)
        
        if not query:
            return jsonify({"error": "查询不能为空"}), 400
        
        if not article_title:
            return jsonify({"error": "文章标题不能为空"}), 400
        
        # 执行检索
        search_results = retriever.search(query, search_type, article_title, k)
        formatted_results = retriever.format_search_results(search_results)
        
        return jsonify({
            "query": query,
            "article_title": article_title,
            "search_type": search_type,
            "results": formatted_results,
            "total_results": len(formatted_results)
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/history', methods=['GET'])
def get_history():
    """获取对话历史"""
    try:
        article_title = request.args.get('article_title', '')
        max_messages = request.args.get('max_messages', type=int, default=0)
        history = chat_manager.get_recent_history(article_title, max_messages)
        
        return jsonify({
            "history": history,
            "summary": chat_manager.get_conversation_summary()
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/history', methods=['DELETE'])
def clear_history():
    """清空对话历史"""
    try:
        article_title = request.args.get('article_title', '')
        chat_manager.clear_history(article_title)
        return jsonify({"message": "对话历史已清空"})
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/documents', methods=['GET'])
def get_documents():
    """获取文档信息"""
    try:
        doc_count = db_manager.get_document_count()
        documents = db_manager.get_all_documents()
        
        # 只返回基本信息，不包含嵌入向量
        doc_info = []
        for doc in documents:
            doc_info.append({
                "title": doc.get("title", ""),
                "category": doc.get("category", ""),
                "content_preview": doc.get("content", "")[:100] + "..." if len(doc.get("content", "")) > 100 else doc.get("content", "")
            })
        
        return jsonify({
            "total_documents": doc_count,
            "documents": doc_info
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/initialize', methods=['POST'])
def initialize_system():
    """初始化系统"""
    try:
        retriever.initialize()
        return jsonify({
            "message": "系统初始化成功",
            "document_count": db_manager.get_document_count(),
            "retriever_ready": retriever.initialized
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    print("🚀 启动RAG智能问答系统...")
    print(f"配置信息:")
    print(f"  - OneAPI地址: {config.ONEAPI_BASE_URL}")
    print(f"  - 数据库: {config.DB_NAME}")
    print(f"  - 最大历史记录: {config.MAX_HISTORY}")
    print(f"  - 流式对话: {'启用' if config.STREAM_ENABLED else '禁用'}")
    
    # 初始化系统
    try:
        retriever.initialize()
        print("✅ 系统初始化完成")
    except Exception as e:
        print(f"❌ 系统初始化失败: {str(e)}")
        print("请检查OneAPI服务和MongoDB连接")
    
    app.run(
        host=config.HOST,
        port=config.PORT,
        debug=config.DEBUG
    ) 